"""
@Project    : crnn
@Module     : json2crop.py
@Author     : wangqinggang@haier.com
@Created    : 2020/11/12 9:34
@Desc       : 对数据进行光照强度的增强，用于数据量少的时候 本地增强文件
"""

from tqdm import tqdm
from glob import glob
import numpy as np
import os
import shutil
import cv2

from expand_sample import lightexchange

def aug_img(src_path, save_path, num_sample=10):
    imgs = glob(src_path + '/*.jpg')

    for index in range(num_sample):
        for im_name in tqdm(imgs):
            im = cv2.imread(im_name)
            name = os.path.basename(im_name)
            new_im = lightexchange(im)
            new_name = os.path.join(save_path, name.replace('.jpg', '')+'_'+str(index)+'.jpg')
            cv2.imwrite(new_name, new_im)
            
def aug_anno(src_path, label_path, save_path, num_sample=10):
    imgs = glob(src_path + '/*.jpg')

    for index in range(num_sample):
        for im_name in tqdm(imgs):
            im = cv2.imread(im_name)
            name = os.path.basename(im_name)
            new_im = lightexchange(im)
            new_name = os.path.join(save_path, name.replace('.jpg', '')+'_'+str(index)+'.jpg')
            cv2.imwrite(new_name, new_im)
            label_name = os.path.join(label_path, name.replace('.jpg', '.xml'))           
            label = os.path.join(save_path, new_name.replace('.jpg', '.xml'))
            shutil.copy(label_name, save_path)
            os.rename(os.path.join(save_path,os.path.basename(label_name)), label)

def aug_mask(src_path, mask_path, save_path, num_sample=10):
    
    imgs = glob(src_path + '/*.jpg')

    for index in range(num_sample):
        for im_name in tqdm(imgs):
            im = cv2.imread(im_name)
            name = os.path.basename(im_name)
            new_im = lightexchange(im)
            new_name = os.path.join(save_path, name.replace('.jpg', '')+'_'+str(index)+'.jpg')
            cv2.imwrite(new_name, new_im)
            label_name = os.path.join(mask_path, name.replace('.jpg', '-label.jpg'))           
            label = os.path.join(save_path, new_name.replace('.jpg', '-label.jpg'))
            shutil.copy(label_name, save_path)
            os.rename(os.path.join(save_path,os.path.basename(label_name)), label)

def segmentation():
    src_path = r'D:\datasets\aotie\ocr\all\origin\c-n'
    anno_path= r'D:\datasets\aotie\ocr\all\origin\xmls'
    save_path = r'D:\datasets\aotie\ocr\all\augment-train'
    aug_anno(src_path, anno_path, save_path, num_sample=3)
#原始代码
def classification_source():
    src_path = r'D:\datasets\haier\tp_classify_data\ori'
    save_path = r'D:\datasets\haier\tp_classify_data\aug'
    dirs = os.listdir(src_path)
    for sub_dir in dirs:
        num_sample = 5
        tmp = os.path.join(src_path, sub_dir)
        imgs = glob(tmp + '/*.jpg')
        save_ = os.path.join(save_path, sub_dir)
        if not os.path.exists(save_):
            os.mkdir(save_)
        for index in range(num_sample):
            for im_name in tqdm(imgs):
                im = cv2.imread(im_name)
                name = os.path.basename(im_name)
                new_im = lightexchange(im)
                new_name = os.path.join(save_, name.replace('.jpg', '')+'_'+str(index)+'.jpg')
                cv2.imwrite(new_name, new_im)
#航空行程单 cpn检测相应修改
def classification():
    src_path = r'D:\PycharmProjects\haier\cpn-det\data\augment_data\img'
    save_img_path = r'D:\PycharmProjects\haier\cpn-det\data\augment_data\aug_img'
    save_txt_path = r'D:\PycharmProjects\haier\cpn-det\data\augment_data\aug_txt'
    num_sample = 100
    imgs = glob(src_path + '/*.jpg')
    for index in range(num_sample):
        for im_name in tqdm(imgs):
            txt_path = im_name.replace('img','txt')[:-4]+'.txt'
            if os.path.isfile(txt_path)==False:
                break
            im = cv2.imread(im_name)
            name = os.path.basename(im_name)
            new_im = lightexchange(im)
            new_img_name = os.path.join(save_img_path, name.replace('.jpg', '').replace('.JPG', '')+'_'+str(index)+'.jpg')
            new_txt_name = os.path.join(save_txt_path, name.replace('.jpg', '').replace('.JPG', '') + '_' + str(index) + '.txt')
            cv2.imwrite(new_img_name, new_im)
            shutil.copyfile(txt_path,new_txt_name)
#航空行程单crnn数据增强
def classification_crnn(f2):
    src_path = r'C:\Users\Administrator\Downloads\1\2\crop_img'
    save_img_path = r'C:\Users\Administrator\Downloads\1\2\aug_crop_img_test'
    num_sample = 3
    #读取原来的txt
    with open(r'C:\Users\Administrator\Downloads\1\2\crop_img_test_random.txt', 'r', encoding='utf-8') as f1:
        lines = f1.readlines()
    for index in range(num_sample):
        count = 0
        for line in tqdm(lines):
            linesplit = line.split('.jpg')
            img_name = str(linesplit[0])+'.jpg'
            label = str(linesplit[-1]).replace(' ','').replace('\n','')

            im = cv2.imread(os.path.join(src_path,img_name))
            new_im = lightexchange(im)
            new_img_name = 'hkcxd_v2_'+str(index)+'_'+str(count)+'.jpg'
            new_img_name_path = os.path.join(save_img_path, new_img_name)
            count+=1
            cv2.imwrite(new_img_name_path, new_im)
            pointstr = new_img_name + ' ' + label.replace('：', ':').replace(' ', '') + '\n'
            f2.write(pointstr)
            # shutil.copyfile(txt_path,new_txt_name)

if __name__ == "__main__":
    # 新生成的txt
    with open(r'C:\Users\Administrator\Downloads\1\2\\' + 'aug_crop_img_test_random.txt', 'w+', encoding='utf-8') as f:
        classification_crnn(f)
    